{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 2 - MultiLayer Split Neural Network\n",
    "\n",
    "<b>Recap:</b> The previous tutorial looked at building a basic SplitNN, where an NN was split into two segments on two seperate hosts. However, there is a lot more that we can do with this technique. An NN can be split any number of times without affecting the accuracy of the model.\n",
    "\n",
    "<b>Description: </b>Here we define a class which can procees a SplitNN of any number of layers. All it needs is a list of distributed models and their optimizers. \n",
    "\n",
    "\n",
    "\n",
    "<img src=\"images/MultiSplitNN.png\" width=\"20%\">\n",
    "\n",
    "\n",
    "In this tutorial, we demonstrate the SplitNN class with a 3 segment distribution [[1](https://arxiv.org/abs/1812.00564)]. This time;\n",
    "\n",
    "- <b>Alice</b>\n",
    "    - Has Model Segment 1\n",
    "    - Has the handwritten images\n",
    "- <b>Bob</b>\n",
    "    - Has model Segment 2\n",
    "- <b>Claire</b> \n",
    "    - Has Model Segment 3\n",
    "    - Has the image labels\n",
    "    \n",
    "We use the exact same model as we used in the previous tutorial, only this time we are splitting over 3 hosts, not two. However, we see the same loss being reported as there is <b>no reduction in accuracy</b> when training in this way. While we only use 3 models this can be done for any arbitrary number of models.\n",
    "\n",
    "Author:\n",
    "- Adam J Hall - Twitter: [@AJH4LL](https://twitter.com/AJH4LL) · GitHub:  [@H4LL](https://github.com/H4LL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch import nn, optim\n",
    "import syft as sy\n",
    "hook = sy.TorchHook(torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SplitNN(torch.nn.Module):\n",
    "    def __init__(self, models, optimizers):\n",
    "        self.models = models\n",
    "        self.optimizers = optimizers\n",
    "        self.outputs = [None]*len(self.models)\n",
    "        self.inputs = [None]*len(self.models)\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        self.inputs[0] = x\n",
    "        self.outputs[0] = self.models[0](self.inputs[0])\n",
    "        \n",
    "        for i in range(1, len(self.models)):\n",
    "            self.inputs[i] = self.outputs[i-1].detach().requires_grad_()\n",
    "            if self.outputs[i-1].location != self.models[i].location:\n",
    "                self.inputs[i] = self.inputs[i].move(self.models[i].location).requires_grad_()               \n",
    "            self.outputs[i] = self.models[i](self.inputs[i])\n",
    "        \n",
    "        return self.outputs[-1]\n",
    "    \n",
    "    def backward(self):\n",
    "        for i in range(len(self.models)-2, -1, -1):\n",
    "            grad_in = self.inputs[i+1].grad.copy()\n",
    "            if self.outputs[i].location != self.inputs[i+1].location:\n",
    "                grad_in = grad_in.move(self.outputs[i].location)\n",
    "            self.outputs[i].backward(grad_in)\n",
    "    \n",
    "    def zero_grads(self):\n",
    "        for opt in self.optimizers:\n",
    "            opt.zero_grad()\n",
    "        \n",
    "    def step(self):\n",
    "        for opt in self.optimizers:\n",
    "            opt.step()\n",
    "    \n",
    "    def train(self):\n",
    "        for model in self.models:\n",
    "            model.train()\n",
    "    \n",
    "    def eval(self):\n",
    "        for model in self.models:\n",
    "            model.eval()\n",
    "            \n",
    "    @property\n",
    "    def location(self):\n",
    "        return self.models[0].location if self.models and len(self.models) else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data preprocessing\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                              transforms.Normalize((0.5,), (0.5,)),\n",
    "                              ])\n",
    "trainset = datasets.MNIST('mnist', download=True, train=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "# Define our model segments\n",
    "\n",
    "input_size = 784\n",
    "hidden_sizes = [128, 640]\n",
    "output_size = 10\n",
    "\n",
    "models = [\n",
    "    nn.Sequential(\n",
    "                nn.Linear(input_size, hidden_sizes[0]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[1], output_size),\n",
    "                nn.LogSoftmax(dim=1)\n",
    "    )\n",
    "]\n",
    "\n",
    "# Create optimisers for each segment and link to them\n",
    "optimizers = [\n",
    "    optim.SGD(model.parameters(), lr=0.03,)\n",
    "    for model in models\n",
    "]\n",
    "\n",
    "# create some workers\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")\n",
    "claire = sy.VirtualWorker(hook, id=\"claire\")\n",
    "\n",
    "# Send Model Segments to model locations\n",
    "model_locations = [alice, bob, claire]\n",
    "for model, location in zip(models, model_locations):\n",
    "    model.send(location)\n",
    "\n",
    "#Instantiate a SpliNN class with our distributed segments and their respective optimizers\n",
    "splitNN =  SplitNN(models, optimizers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(x, target, splitNN):\n",
    "    \n",
    "    #1) Zero our grads\n",
    "    splitNN.zero_grads()\n",
    "    \n",
    "    #2) Make a prediction\n",
    "    pred = splitNN.forward(x)\n",
    "  \n",
    "    #3) Figure out how much we missed by\n",
    "    criterion = nn.NLLLoss()\n",
    "    loss = criterion(pred, target)\n",
    "    \n",
    "    #4) Backprop the loss on the end layer\n",
    "    loss.backward()\n",
    "    \n",
    "    #5) Feed Gradients backward through the nework\n",
    "    splitNN.backward()\n",
    "    \n",
    "    #6) Change the weights\n",
    "    splitNN.step()\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(epochs):\n",
    "    running_loss = 0\n",
    "    splitNN.train()\n",
    "    for images, labels in trainloader:\n",
    "        images = images.send(models[0].location)\n",
    "        images = images.view(images.shape[0], -1)\n",
    "        labels = labels.send(models[-1].location)\n",
    "        loss = train(images, labels, splitNN)\n",
    "        running_loss += loss.get()\n",
    "\n",
    "    else:\n",
    "        print(\"Epoch {} - Training loss: {}\".format(i, running_loss/len(trainloader)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, dataloader, dataset_name):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in dataloader:\n",
    "            data = data.view(data.shape[0], -1).send(model.location)\n",
    "            output = model(data).get()\n",
    "            pred = output.data.max(1, keepdim=True)[1]\n",
    "            correct += pred.eq(target.data.view_as(pred)).sum()\n",
    "\n",
    "    print(\"{}: Accuracy {}/{} ({:.0f}%)\".format(dataset_name, \n",
    "                                                correct,\n",
    "                                                len(dataloader.dataset), \n",
    "                                                100. * correct / len(dataloader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testset = datasets.MNIST('mnist', download=True, train=False, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)\n",
    "test(splitNN, testloader, \"Test set\")\n",
    "test(splitNN, trainloader, \"Train set\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
